#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <arpa/inet.h>

#include <boost/coroutine2/all.hpp>
#include <iostream>
#include <thread>
#include <chrono>
#include <string>

#include "executor.hpp"
#include "asyncio.hpp"

#define THREAD_ID   " (thread id: " << std::hex << std::this_thread::get_id() << std::dec << ")"

static void spawn(coro::Executor &executor,
                  const std::string &name,
                  int cost_time, int val)
{
    executor.spawn<int>([&name, cost_time, val](coro::Context<int> &ctx) {
        std::string CORO = name + " >>>> ";
        std::cout << CORO << "Start" << THREAD_ID << std::endl;

        coro::Waker<int> waker = ctx.waker();
        std::thread t{[waker, &CORO, cost_time, val]() mutable {
            using namespace std::chrono_literals;

            std::cout << CORO << "Running an async task for "
                      << cost_time << "s" << THREAD_ID << std::endl;
            std::this_thread::sleep_for(std::chrono::seconds(cost_time));
            std::cout << CORO << "Async task done with value " << val << THREAD_ID << std::endl;
            waker.wake(val);
        }};
        t.detach();

        std::cout << CORO << "Wait for async task to finish..." << THREAD_ID << std::endl;
        auto v = ctx.yield();
        std::cout << CORO << "Async task finished!" << THREAD_ID << std::endl;
        if (v) {
            std::cout << CORO << "Async task result: " << v.value() << THREAD_ID << std::endl;
        }
    });
}

int demo_thread_waker(int argc, char *argv[])
{
    std::cout << "main >>>> " << "Start" << THREAD_ID << std::endl;

    coro::Executor executor;

    spawn(executor, "coro 1", 4, 111);
    spawn(executor, "coro 2", 2, 222);

    executor.run();
    return 0;
}

// TODO: fd for regular file is not supported by epoll
int demo_zero_to_null(int argc, char *argv[])
{
    coro::Executor executor;

    for (int i = 0; i < 9; ++ i) {
        executor.spawn<asyncio::IoResult>([i](coro::Context<asyncio::IoResult> &ctx) {
            auto CORO = (boost::format("CORO %1% >>> ") % i).str();
            const int BUFF_SIZE = 128;
            char buf[BUFF_SIZE] = {0};

            int dev_zero = open("/dev/zero", O_CLOEXEC|O_NONBLOCK);
            if (dev_zero < 0) {
                std::cerr << CORO << "failed to open /dev/zero: "
                          << utils::strerror(errno) << std::endl;
                return;
            }
            int dev_null = open("/dev/null", O_CLOEXEC|O_NONBLOCK);
            if (dev_null < 0) {
                close(dev_zero);
                std::cerr << CORO << "failed to open /dev/zero: "
                          << utils::strerror(errno) << std::endl;
                return;
            }

            while (true) {
                ssize_t res = asyncio::read_async(ctx, dev_zero, buf, BUFF_SIZE);
                if (res < 0) {
                    std::cerr << CORO << "failed to read /dev/zero: "
                              << utils::strerror(errno) << std::endl;
                    break;
                } else {
                    std::cerr << CORO << "read " << res
                              << "bytes from /dev/zero" << std::endl;
                }
                res = asyncio::write_async(ctx, dev_null, buf, res);
                if (res < 0) {
                    std::cerr << CORO << "failed to write /dev/null: "
                              << utils::strerror(errno) << std::endl;
                    break;
                } else {
                    std::cerr << CORO << "write " << res
                              << "bytes to /dev/null" << std::endl;
                }
            }
            close(dev_zero);
            close(dev_null);
        });
    }
    executor.run();
    return 0;
}

void demo_sock_server_handle_sock(coro::Context<asyncio::IoResult> &ctx,
                                  int peer_fd, sockaddr_in peer_addr)
{
    char addr[24] = {0};
    if (!inet_ntop(AF_INET, (void *)&peer_addr.sin_addr.s_addr, addr, 24)) {
        std::cerr << "failed to parse ip: " << utils::strerror(errno) << std::endl;
        return;
    }
    int port = ntohs(peer_addr.sin_port);
    std::string CORO = (boost::format("%1%:%2% >>> ") % addr % port).str();
    std::cerr << CORO << "start to handle data" << std::endl;

    const int BUFF_SIZE = 4096;
    char buff[BUFF_SIZE] = {0};
    while (true) {
        ssize_t r = asyncio::read_async(ctx, peer_fd, buff, BUFF_SIZE);
        if (r < 0) {
            std::cerr << CORO << "failed to read from socket: "
                      << utils::strerror(errno) << std::endl;
            break;
        }
        std::cerr << CORO << "read " << r << " bytes from socket" << std::endl;
        char *p = buff;
        while (r > 0) {
            ssize_t w = asyncio::write_async(ctx, peer_fd, p, r);
            if (w < 0) {
                std::cerr << CORO << "failed to write to socket: "
                          << utils::strerror(errno) << std::endl;
                r = -errno;
                break;
            }
            std::cerr << CORO << "write " << r << " bytes to socket" << std::endl;
            r -= w;
            p += w;
        }
        if (r < 0) {
            break;
        }
    }
    close(peer_fd);
}

// CMD socket_server IP PORT
int demo_socket_server(int argc, char *argv[])
{
    if (argc < 4) {
        std::cerr << "Too few arguments" << std::endl;
        return EINVAL;
    }
    const char *ip = argv[2];
    int port = 0;
    try {
        port = std::stoi(argv[3]);
    } catch (const std::exception &e) {
        std::cerr << "failed to get port: " << e.what() << std::endl;
        return EINVAL;
    }

    coro::Executor executor;

    executor.spawn<asyncio::IoResult>([ip, port](coro::Context<asyncio::IoResult> &ctx) {
        int r = 0;
        int srv_fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
        if (srv_fd < 0) {
            std::cerr << "failed to create server socket: "
                    << utils::strerror(errno) << std::endl;
            return;
        }
        sockaddr_in addr = {
            .sin_family = AF_INET,
            .sin_port = htons(uint16_t(port)),
            .sin_addr = {
                .s_addr = inet_addr(ip)
            }
        };
        if (addr.sin_addr.s_addr == INADDR_NONE) {
            std::cerr << "wrong address: " << ip << std::endl;
            close(srv_fd);
            return;
        }
        r = bind(srv_fd, (struct sockaddr *)&addr, sizeof(addr));
        if (r < 0) {
            std::cerr << "failed to bind: " << utils::strerror(errno) << std::endl;
            close(srv_fd);
            return;
        }
        r = listen(srv_fd, 256);
        if (r < 0) {
            std::cerr << "failed to listen: " << utils::strerror(errno) << std::endl;
            close(srv_fd);
            return;
        }
        while (true) {
            sockaddr_in peer_addr = {0};
            socklen_t peer_addr_len = sizeof(peer_addr);
            int peer_sock = asyncio::accept_async(ctx, srv_fd,
                                (sockaddr *)&peer_addr, &peer_addr_len);
            if (peer_sock < 0) {
                std::cerr << "failed to accepth sock: "
                          << utils::strerror(errno) << std::endl;
                close(srv_fd);
                break;
            }
            ctx.get_executor()->spawn<asyncio::IoResult>(
                [peer_sock, peer_addr](coro::Context<asyncio::IoResult> &ctx) {
                    demo_sock_server_handle_sock(ctx, peer_sock, peer_addr);
                });
        }
    });

    executor.run();
    return 0;
}

void demo_socket_client_connect_server(coro::Context<asyncio::IoResult> &ctx,
                                       const char *ip, int port, int no)
{
    const int BUFFSIZE = 4096;
    std::vector<char> buff;
    buff.reserve(BUFFSIZE);
    const char *p = buff.data();
    int r = 0;

    std::string CORO = (boost::format("CONN %1% >>> ") % no).str();

    int fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
    if (fd < 0) {
        std::cerr << CORO << "failed to create socket: "
                  << utils::strerror(errno) << std::endl;
        return;
    }
    sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_port = htons(uint16_t(port)),
        .sin_addr = {
            .s_addr = inet_addr(ip)
        }
    };
    if (addr.sin_addr.s_addr == INADDR_NONE) {
        std::cerr << CORO << "wrong addr: " << ip << std::endl;
        close(fd);
        return;
    }

    r = asyncio::connect_async(ctx, fd, (sockaddr *)&addr, sizeof(addr));
    if (r < 0) {
        std::cerr << CORO << "failed to connect to server " << ip << ":" << port << " "
                  << utils::strerror(errno) << std::endl;
        close(fd);
        return;
    }

    // Spawn a coroutine for reading
    int read_fd = dup(fd);
    if (read_fd < 0) {
        std::cerr << CORO << "failed to dup fd: "
                  << utils::strerror(errno) << std::endl;
        close(fd);
        return;
    }
    ctx.get_executor()->spawn<asyncio::IoResult>(
        [read_fd, &CORO, &BUFFSIZE](coro::Context<asyncio::IoResult> &ctx) {
            char *buff = new char[BUFFSIZE];
            int r = 0;

            while (true) {
                r = asyncio::read_async(ctx, read_fd, buff, BUFFSIZE);
                if (r < 0) {
                    std::cerr << CORO << "failed to read from socket: "
                              << utils::strerror(errno) << std::endl;
                    break;
                }
                std::cerr << CORO << "read " << r << " bytes from socket"
                          << std::endl;
            }
            delete [] buff;
            close(read_fd);
        });

    // Current coroutine is used for write
    while (true) {
        r = asyncio::write_async(ctx, fd, p, BUFFSIZE);
        if (r < 0) {
            std::cerr << CORO << "failed to write to socket: "
                      << utils::strerror(errno) << std::endl;
            break;
        }
        std::cerr << CORO << "write " << r << " bytes to socket" << std::endl;
    }
    close(fd);
}

// CMD socket client IP PORT CLIENT_COUNT
int demo_socket_client(int argc, char *argv[])
{
    if (argc < 4) {
        std::cerr << "Too few arguments" << std::endl;
        return EINVAL;
    }
    const char *ip = argv[2];
    int port = 0;
    int client_count = 0;
    try {
        port = std::stoi(argv[3]);
        client_count = std::stoi(argv[4]);
    } catch (const std::exception &e) {
        std::cerr << "failed to parse port or client_count: " << e.what() << std::endl;
        return EINVAL;
    }

    coro::Executor executor;

    for (int i = 0; i < client_count; ++ i) {
        executor.spawn<asyncio::IoResult>(
            [ip, port, i](coro::Context<asyncio::IoResult> &ctx) {
                demo_socket_client_connect_server(ctx, ip, port, i);
            });
    }

    executor.run();

    return 0;
}

struct {
    const char *name;
    int (*entry)(int argc, char *argv[]);
} demos[] = {
    {
        "thread_waker",
        demo_thread_waker
    },
    {
        "zero_to_null",
        demo_zero_to_null
    },
    {
        "socket_server",
        demo_socket_server
    },
    {
        "socket_client",
        demo_socket_client
    },
    {nullptr, nullptr}
};

int main(int argc, char *argv[])
{
    int res = 0;

    if (argc > 1) {
        std::string demo_name = argv[1];
        int i = 0;
        for (i = 0; demos[i].name != nullptr; ++ i) {
            if (demos[i].name == demo_name) {
                res = demos[i].entry(argc, argv);
                break;
            }
        }
        if (demos[i].name == nullptr) {
            std::cerr << "unknown demo name: " << demo_name << std::endl;
            res = EINVAL;
        }
    } else {
        std::cerr << "no enough argument" << std::endl;
        res = EINVAL;
    }

    return res;
}
